import subprocess as sp

import os
import time
from plot import AgentPloter

# gpu_clock_list=[180000000, 267000000, 355000000, 430000000]
# dir_thermal='/sys/devices/virtual/thermal'

class GPU:
    # The number of GPUs on the device
    n_gpu = 1
    # The temperature info collected during runtime
    temp_data = []

    ########################################
    # GPU instances initialization methods #
    ########################################
    def __init__(self, ip: str):
        self.ip        = ip
        self.freq_data = []

        self.freq_list    = self.init_available_freq()
        self.max_freq_idx = len(self.freq_list) - 1
        self.freq_idx     = self.max_freq_idx

        self.init_scaling_range() # TODO: not always minimum
    
    def init_available_freq(self):
        '''
        Initialize the available frequency list of current GPU instance
        '''
        fname  = '/sys/class/kgsl/kgsl-3d0/devfreq/available_frequencies'
        output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
        print(f'[gpu] available frequencies: {output.decode("utf-8").strip()}')
        return [int(f) for f in output.split()]

    def init_scaling_range(self):
        '''
        Initialize the frequency scaling range of current GPU instance
        '''
        # Set scaling min frequency
        fname = '/sys/class/kgsl/kgsl-3d0/devfreq/min_freq'
        sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"echo', str(self.freq_list[0])+' >', fname+'\"'])
        output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
        min_freq = int(output.decode('utf-8').strip())
        # Set scaling max frequency
        fname = '/sys/class/kgsl/kgsl-3d0/devfreq/max_freq'
        sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"echo', str(self.freq_list[self.max_freq_idx])+' >', fname+'\"'])
        output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
        max_freq = int(output.decode('utf-8').strip())
        # Print prompt
        # print(f'[gpu] scaling from {min_freq} to {max_freq}')

    #####################################
    # GPU instances information methods #
    #####################################
    def set_freq(self, freq_idx: int):
        '''
        Set the frequency of current GPU instance
        '''
        # TODO：调成level
        self.freq_idx = freq_idx
        fname = '/sys/class/kgsl/kgsl-3d0/devfreq/userspace/set_freq'
        sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"echo', str(self.freq_list[freq_idx])+' >', fname+'\"'])

    def log_freq(self):
        '''
        Log the frequency of current GPU instance
        '''
        fname  = '/sys/class/kgsl/kgsl-3d0/devfreq/cur_freq'
        output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
        freq   = int(output.decode('utf-8').strip()) / 1e6
        self.freq_data.append(freq)

    def get_freq(self):
        '''
        Get the frequency of current GPU instance
        '''
        # TODO: 调成level
        fname  = '/sys/class/kgsl/kgsl-3d0/devfreq/cur_freq'
        output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
        freq   = int(output.decode('utf-8').strip()) / 1e6
        return freq

    ###############################
    # Static methods for all GPUs #
    ###############################
    @staticmethod
    def set_scaling_governor(ip: str, gov='msm-adreno-tz'):
        '''
        Set all GPUs' scaling governors to gov
            options: msm-adreno-tz, userspace, powersave, performance, simple_ondemand, ...
        '''
        # Check if the input gov is legal
        fname  = '/sys/class/kgsl/kgsl-3d0/devfreq/available_governors'
        output = sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"cat', fname+'\"'])
        avail_govs = output.decode('utf-8').split()
        if gov not in avail_govs:
            raise ValueError(f'governor must be one of {avail_govs}')
        # Set the governor
        fname = '/sys/class/kgsl/kgsl-3d0/devfreq/governor'
        sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"echo', gov+' >', fname+'\"'])
        # Print prompt
        output = sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"cat', fname+'\"'])
        print(f'[GPU] set the governor to {output.decode("utf-8").strip()}')

    @staticmethod
    def get_overall_temp(ip: str):
        '''
        Get the temperature
        '''
        # TODO: thermal file
        # 输出所有的termal_zone的温度来比对
        fname  = '/sys/devices/virtual/thermal/thermal_zone10/temp'
        output = sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"cat', fname+'\"'])
        temp   = int(output.decode('utf-8').strip()) / 1e3
        print(f'[GPU] temperature: {temp}')
        GPU.temp_data.append(temp)

if __name__ == '__main__':
    # 有线连接
    title = ['min_freq', 'max_freq', 'cur_freq', 'gpu_load']
    ploter = AgentPloter(4, 1, 4, title)
    dir = '/sys/class/kgsl/kgsl-3d0/devfreq/'
    
    nfile = 4
    freq_data = [[] for _ in range(nfile)]

    for t in range(20):
        # 修改min max
        if (t == 3):
            # 修改min_freq
            fname = dir + title[0]
            sp.check_output(['adb', 'shell', 'su -c', '\"echo', str(585000000)+' >', fname+'\"'])

        if (t == 15):
            # 复原min_freq
            fname = dir + title[0]
            sp.check_output(['adb', 'shell', 'su -c', '\"echo', str(257000000)+' >', fname+'\"'])

        # 获取GPU info
        for i in range(nfile):
            fname = dir + title[i]
            output = sp.check_output(['adb', 'shell', 'su -c', '\"cat', fname+'\"'])
            freq_data[i].append(int(output.decode('utf-8').strip()))

        # 画图
        ploter.update(freq_data[0], freq_data[1], freq_data[2], freq_data[3])

        time.sleep(2)
    
    os.system('pause')